

#include "FixedSizeAllocator.h"
#include <stdlib.h>
#include <string.h>
#include <assert.h>
#include <sstream>

namespace FlameEngine
{
	////////////////////////////////////////////////////////////////
	MemChunk::MemChunk(rsize_t blockSize, uint8_t count)
	{
		m_pData = new uint8_t[blockSize * count];
		memset(m_pData, 0, blockSize * count);
		m_firstFreeBlock = 0;
		m_freeBlockCount = count;

		uint8_t i = 0;
		uint8_t* p = m_pData;
		for(; i < m_freeBlockCount; p += blockSize)
		{
			*p = ++i;
		}
	}

	MemChunk::~MemChunk()
	{
		if(m_pData)
			delete[] m_pData;

		m_pData = nullptr;
	}

	void* MemChunk::Allocate(size_t blockSize)
	{
		if (m_freeBlockCount == 0)
			return 0;

		uint8_t* pResult = m_pData + m_firstFreeBlock * blockSize;
		m_firstFreeBlock = *pResult;
		m_freeBlockCount--;

		return pResult;
	}

	void MemChunk::Deallocate(void* p, size_t blockSize)
	{
		assert(p >= m_pData);
		uint8_t* pBlock2Release = static_cast<uint8_t*>(p);

		//alignment check
		assert((pBlock2Release - m_pData) % blockSize == 0);
		memset(pBlock2Release, 0, blockSize);
		*pBlock2Release = m_firstFreeBlock;
		m_firstFreeBlock = static_cast<uint8_t>((pBlock2Release - m_pData) / blockSize);
		m_freeBlockCount++;
	}

	////////////////////////////////////////////////////////////////
	FixedSizeAllocator::FixedSizeAllocator(size_t blockSize, uint8_t maxBlockPerChunk)
		:m_blockSize(blockSize),
		m_maxBlockPerChunk(maxBlockPerChunk),
		m_chunkSize(blockSize * maxBlockPerChunk),
		m_pAllocChunk(nullptr),
		m_pDeallocChunk(nullptr),
		m_totalEmptyBlockCount(0)
	{
		m_threadId = std::this_thread::get_id();
	}

	FixedSizeAllocator::~FixedSizeAllocator()
	{
		for(size_t i = 0; i < m_memChunks.size(); i++)
			delete m_memChunks[i];

		m_pAllocChunk = nullptr;
		m_pDeallocChunk = nullptr;
	}

	void* FixedSizeAllocator::Allocate()
	{
		assert(m_threadId == std::this_thread::get_id() && "use FixedSizeAllocator in different thread");
		return AllocateImp();
	}

	void FixedSizeAllocator::Deallocate(void* p)
	{
		assert(m_threadId == std::this_thread::get_id() && "use FixedSizeAllocator in different thread");
		return DeallocateImp(p);
	}

	void* FixedSizeAllocator::AllocateImp()
	{
		//check if we have free space in current chunk
		if (m_pAllocChunk == nullptr || m_pAllocChunk->GetFreeCount() == 0)
		{
			m_pAllocChunk = nullptr;
			//if not find/create one
			for (size_t i = 0; i < m_memChunks.size(); i++)
			{
				if (m_memChunks[i]->GetFreeCount() > 0)
				{
					m_pAllocChunk = m_memChunks[i];
					break;
				}
			}

			if (m_pAllocChunk == nullptr)
			{
				m_pAllocChunk = new MemChunk(m_blockSize, m_maxBlockPerChunk);
				m_memChunks.push_back(m_pAllocChunk);
				m_totalEmptyBlockCount += m_maxBlockPerChunk;
				if (m_pDeallocChunk == nullptr)
					m_pDeallocChunk = m_pAllocChunk;
			}
		}

		assert(m_pAllocChunk != nullptr && "memory chunk is null!!");
		assert(m_pAllocChunk->GetFreeCount() > 0 && "memory chunk is full");
		assert(m_totalEmptyBlockCount > 0 && "total empty block number is wrong");

		m_totalEmptyBlockCount--;
		return m_pAllocChunk->Allocate(m_blockSize);
	}

	void FixedSizeAllocator::DeallocateImp(void* p)
	{
		assert(p != nullptr && "deallocate null pointer");
		assert(m_pDeallocChunk != nullptr && "deallocate chunk is null");

		if (!m_pDeallocChunk->HasPointer(p, m_chunkSize))
		{
			for (size_t i = 0; i < m_memChunks.size(); i++)
			{
				if (m_memChunks[i]->HasPointer(p, m_chunkSize))
				{
					m_pDeallocChunk = m_memChunks[i];
					break;
				}
			}
		}

		m_pDeallocChunk->Deallocate(p, m_blockSize);
		m_totalEmptyBlockCount++;

		//check if we can free any empty chunk
		if (m_totalEmptyBlockCount >= m_maxBlockPerChunk * 2)
		{
			for (size_t i = 0; i < m_memChunks.size(); i++)
			{
				if (m_memChunks[i]->GetFreeCount() == m_maxBlockPerChunk)
				{
					//delete chunk
					delete m_memChunks[i];
					m_memChunks[i] = m_memChunks.back();
					m_memChunks.pop_back();
					m_totalEmptyBlockCount -= m_maxBlockPerChunk;

					m_pAllocChunk = nullptr;
					m_pDeallocChunk = nullptr;
					//update current chunk pointer
					for (size_t j = 0; j < m_memChunks.size(); j++)
					{
						if (m_memChunks[i]->GetFreeCount() > 0)
						{
							m_pAllocChunk = m_memChunks[i];
							m_pDeallocChunk = m_pDeallocChunk;
							break;
						}
					}

					break;
				}
			}
		}
	}

	////////////////////////////////////////////////////////////////
	FixedSizeAllocatorThreadSafe::FixedSizeAllocatorThreadSafe(size_t blockSize, uint8_t maxBlockPerChunk)
		:FixedSizeAllocator(blockSize, maxBlockPerChunk)
	{
	}

	FixedSizeAllocatorThreadSafe::~FixedSizeAllocatorThreadSafe()
	{
	}

	void* FixedSizeAllocatorThreadSafe::Allocate()
	{
		std::lock_guard<std::mutex> lock(m_mutex);
		return FixedSizeAllocator::AllocateImp();
	}

	void FixedSizeAllocatorThreadSafe::Deallocate(void* p)
	{
		std::lock_guard<std::mutex> lock(m_mutex);
		FixedSizeAllocator::DeallocateImp(p);
	}
}